import math
import numpy as np
from scipy.signal import find_peaks
import statsmodels.api as sm
# from draw_pictures import *

# downsampling_ratio = 2
# second = 10
# framerate = 2000

ratio_ = 1e9

# 定义一个函数，实现对音频在特定位置处的切割
def find_envelope(wave_data,downsampling_ratio,framerate):
    '''
    :param downsampling_ratio: 重采样比率
    :param wave_data: 预处理后的音频数据
    :return: 包络时间轴, 包络值
    '''

    # 开始计算香农能量 ------------------------------
    timesE = []  # 对应的时间戳
    E = []  # 对应的香农熵
    t = []  # 对应的分段次数
    k = 1

    seg_length = 16
    overlap_seg_length = 8
    order = 3

    while (k - 1) * (seg_length - overlap_seg_length) + seg_length - 1 < len(wave_data):
        res = 0
        for j in range((k - 1) * (seg_length - overlap_seg_length),
                       (k - 1) * (seg_length - overlap_seg_length) + seg_length):
            x = (abs(wave_data[j])) ** order
            x = x + 1e-11
            y = x * math.log(x)
            res = res + y

        res = res * -1 / seg_length
        E.append(res)
        timesE.append(((k - 1) * (seg_length - overlap_seg_length) + seg_length / 2 - 1) * downsampling_ratio / (framerate))
        t.append(k)
        k = k + 1

    timesE = np.array(timesE)
    E = np.array(E)
    E_mean = E.mean()
    E_std = E.std()
    E = (E - E_mean) / E_std  # 标准化处理

    # MaxE = float(max(E))
    # MinE = float(min(E))
    # E = np.array([float(k - MinE) / (MaxE - MinE) for k in E])
    # E = E * 5
    maxv = max(abs(E))
    if maxv == -min(E):
        E = abs(E)
    # 平滑处理 ------------------------------
    lowess = sm.nonparametric.lowess
    resE = lowess(E, timesE, frac=10.0*downsampling_ratio/framerate, it=4, delta=8.0*downsampling_ratio/framerate)
    # resE = lowess(E, timesE, frac=10.0*downsampling_ratio/framerate, it=3, delta=8.0*downsampling_ratio/framerate)

    # drawEnvelopeCurve(timesE,E,resE[:,1])

    return timesE,resE[:,1]

def find(x, v, threshold2, window_width,timesE,envelope):
    '''
    :param x: 大于上阈值的极值点索引
    :param v: v∈[1,2]
    :return: 边界的包络索引
    '''

    minx = x
    minv = 1e9

    if v == 1: # 搜索左边界
        t = x
        while t >= 0 and timesE[x]-timesE[t] <= window_width:
            if envelope[t] <= threshold2 and ( envelope[t] < minv or minv==1e9 ):
                minx = t
                minv = envelope[t]

            t = t - 1

        return minx

    elif v == 2: # 搜索右边界
        t = x
        while t < timesE.shape[0] and timesE[t] - timesE[x] <= window_width:

            if envelope [t] <= threshold2 and ( envelope[t] < minv or minv==1e9 ):
                minx = t
                minv = envelope[t]

            t = t + 1

        return minx

def merge_interval(interval,maximum):
    flag = False
    last = -2e9
    for seg in interval:
        l = seg[0]
        r = seg[1]

        if last == -2e9:
            last = r
        else:
            if l <= last:
                flag = True
                break
            last = r
    if flag == True:
        # 区间合并（可能会存在重叠区间）

        interval = sorted(interval)

        res_itv = []  # 存储最后合并的去间，每个item都是一个list [l,r]表示一个区间
        res_maximum = []

        start = -2e9
        end = -2e9
        maxv = -1e9

        for i in range(len(interval)):
            seg = interval[i]
            l = seg[0]
            r = seg[1]
            v = maximum[i]

            if end < l:
                if start != -2e9:
                    res_itv.append([start, end])
                    res_maximum.append(maxv)
                    maxv = v

                if start == -2e9:
                    maxv = v

                start = l
                end = r
            else:
                end = max(end, r)
                maxv = max(maxv, v)

        if start != -2e9:
            res_itv.append([start, end])
            res_maximum.append(maxv)

        interval = res_itv
        del res_itv

        maximum = res_maximum
        del res_maximum

    return interval,maximum

def search_time_door(timesE,envelope,peaks,second):
    '''
    :param time:
    :param envelope:
    :return:
    '''

    if len(peaks) == 0:
        interval = []
        maximum = []
        threshold1 = 0
        threshold2 = 0
    else:
        envelope_values = [envelope[x] for x in peaks]
        envelope_values = sorted(envelope_values)

        global ratio_
        if ratio_ == 1e9:
            ratio_ = 0.005888 * (second**2) - 0.3153 * second + 6.109
        idx_ratio = 1 - second * ratio_/ len(peaks)
        # print("idx_ratio ", idx_ratio)

        index1 = idx_ratio * len(envelope_values)
        index2 = 0.85 * len(envelope_values)
        if index1 == index1:
            index1 = int(index1)
        else:
            index1 = 0

        if index2 == index2:
            index2 = int(index2)
        else:
            index2 = 0

        threshold1 = envelope_values[index1]
        threshold2 = envelope_values[index2]
        window_width = 0.08

        interval = []  # 记录由每个极值点搜出的区间
        maximum = []  # 记录每个区间的极大值

        for i in range(len(peaks)):
            # x 表示: 极值在E(包络)中的索引, 例如第136个包络值
            x = peaks[i]
            y = envelope[x]

            if y > threshold1:
                lmin = find(x,1,threshold2,window_width,timesE,envelope)
                rmin = find(x,2,threshold2,window_width,timesE,envelope)

                subinterval = [lmin, rmin]
                interval.append(subinterval)
                maximum.append(y)

        interval, maximum = merge_interval(interval,maximum)

    return interval, maximum, threshold1, threshold2

def sort_interval(original_interval, new_interval, original_maximum, new_maximum):

    tn_interval = []
    length = len(new_interval)

    for i in range(length):
        v = new_interval[i]
        v.append(new_maximum[i])
        tn_interval.append(v)

    tn_interval = sorted(tn_interval)
    new_interval = []
    new_maximum = []

    for i in range(length):
        itv = tn_interval[i][0:2]
        maxv = tn_interval[i][2]
        new_interval.append(itv)
        new_maximum.append(maxv)

    tmp_interval = []
    tmp_maximum = []
    last = 0
    kidx = 0
    for midx in range(len(new_interval)):
        cur_l = new_interval[midx][0]
        while kidx + 1 < len(original_interval) and original_interval[kidx][0] < cur_l and original_interval[kidx + 1][0] < cur_l:
            kidx += 1

        tmp_interval.extend(original_interval[last:kidx + 1])
        tmp_maximum.extend(original_maximum[last:kidx + 1])
        tmp_interval.append(new_interval[midx])
        tmp_maximum.append(new_maximum[midx])
        last = kidx + 1

    if kidx < len(original_interval):
        tmp_interval.extend(original_interval[last:])
        tmp_maximum.extend(original_maximum[last:])

    return tmp_interval, tmp_maximum

def filter_interval(timesE,interval,maximum,envelope,peaks):
    dist = []
    for i in range(len(interval) - 1):
        r = timesE[interval[i][1]]
        nl = timesE[interval[i + 1][0]]

        dist.append(nl - r)

    sorted_dist = sorted(dist)
    y_min = min([envelope[x] for x in peaks]) - 0.1
    is_delete1 = []
    is_merge = [0 for i in range(len(interval))]

    th = 0.6
    for i in range(len(interval)-1):
        r = timesE[interval[i][1]]
        nl = timesE[interval[i + 1][0]]
        d = nl - r

        if d < 0.047:
            h1 = maximum[i] - y_min
            h2 = maximum[i+1] - y_min

            if h1 < h2 * th:
                is_delete1.append(i)
            elif h2 < h1 * th:
                is_delete1.append(i+1)
            elif (h2 > h1 and h1 >= h2 * th) or (h1 > h2 and h2 >= h1 * th):
                is_merge[i] = 1
                is_merge[i+1] = 2

    if len(is_delete1) != 0 or np.sum(is_merge) != 0:
        res_interval = []
        res_maximum = []

        for i in range(len(interval)):
            if is_delete1.count(i):
                continue

            if is_merge[i] == 1 and is_merge[i+1] == 2:
                l = interval[i][0]
                r = interval[i+1][1]
                maxh = max(maximum[i],maximum[i+1])
                res_interval.append([l,r])
                res_maximum.append(maxh)
                continue

            if is_merge[i] == 2:
                continue

            res_interval.append(interval[i])
            res_maximum.append(maximum[i])

        interval = res_interval
        maximum = res_maximum
        del res_maximum
        del res_interval

    dist2 = []
    for i in range(len(interval) - 1):
        r = timesE[interval[i][1]]
        nl = timesE[interval[i + 1][0]]

        dist2.append(nl - r)

    sorted_dist2 = sorted(dist2)

    # x1 = 0
    # x2 = 0
    # var_sum = 1e9
    #
    # for i in range(len(sorted_dist2) - 1):
    #     # if i + 1 == 1:
    #     #     continue
    #     for j in np.arange(i + 1, len(sorted_dist2)):
    #         # if j - i == 1 or len(sorted_dist) - j - 1 == 1:
    #         #     continue
    #
    #         seg1 = sorted_dist2[0:i + 1]
    #         seg2 = sorted_dist2[i + 1:j + 1]
    #         seg3 = sorted_dist2[j + 1:]
    #
    #         tvar1 = np.var(seg1)
    #         tvar2 = np.var(seg2)
    #         tvar3 = np.var(seg3)
    #
    #         tvar_sum = tvar1 + tvar2 + tvar3
    #         if var_sum > tvar_sum:
    #             var_sum = tvar_sum
    #             x1 = i
    #             x2 = j
    #
    # threshold = sorted_dist2[x1]

    threshold = 0.07

    is_delete2 = []
    for seg_idx in range(1, len(interval) - 1):
        ldist = dist2[seg_idx - 1]
        rdist = dist2[seg_idx]
        if ldist < threshold and rdist < threshold:
            is_delete2.append(seg_idx)

    if len(is_delete2) != 0:
        res_interval = []
        res_maximum = []

        for idx in range(len(interval)):
            if is_delete2.count(idx):
                continue
            res_interval.append(interval[idx])
            res_maximum.append(maximum[idx])

        interval = res_interval
        maximum = res_maximum

    ## 第二种筛选
    ndist = []
    for i in range(len(interval) - 1):
        r = timesE[interval[i][1]]
        nl = timesE[interval[i + 1][0]]

        ndist.append(nl - r)

    sorted_ndist = sorted(ndist)

    # print("ndist ",len(ndist))
    # print(ndist)

    x = 0
    var_sum = 1e9

    for i in range(len(sorted_ndist) - 1):
        seg1 = sorted_ndist[0:i + 1]
        seg2 = sorted_ndist[i + 1:]

        tvar1 = np.var(seg1)
        tvar2 = np.var(seg2)

        tvar_sum = tvar1 + tvar2
        if var_sum > tvar_sum:
            var_sum = tvar_sum
            x = i

    systolic_seg = sorted_ndist[0:x + 1]
    if len(systolic_seg) != 0:
        systolic_max = max(systolic_seg)

        is_delete3 = []
        for idx in range(1, len(interval) - 1):
            if is_delete3.count(idx):
                continue
            d1 = ndist[idx - 1]
            d2 = ndist[idx]

            if d1 <= systolic_max and d2 <= systolic_max:
                h1 = maximum[idx - 1]
                h2 = maximum[idx + 1]
                h = maximum[idx]
                if h1 < h2 and h1 < h * 0.1:
                    is_delete2.append(idx - 1)
                elif h1 > h2 and h2 < h * 0.1:
                    is_delete2.append(idx + 1)

        res_interval = []
        res_maximum = []

        for idx in range(0, len(interval)):
            if is_delete3.count(idx):
                continue
            res_interval.append(interval[idx])
            res_maximum.append(maximum[idx])

        interval = res_interval
        maximum = res_maximum

        del res_maximum
        del res_interval

    return interval, maximum

def envelope_segmentation(timesE, envelope, second):
    '''
    :param timesE: 包络时间轴
    :param envelope: 包络值
    :param second: 读取音频秒数
    :return:
    '''

    # 找极值点 + 绘图 ------------------------------------
    peaks, _ = find_peaks(envelope, distance=1) # peaks: type -> list , 一组极值点的索引
    # drawPeaks(timesE,envelope,peaks)

    sec = second
    interval, maximum, threshold1, threshold2 = search_time_door(timesE,envelope,peaks,sec)

    # drawTimeDoor(timesE,envelope,peaks,second,threshold1,threshold2,True,interval,maximum)

    for epoch in range(5):
        dist = []
        for i in range(len(interval)-1):
            r = timesE[interval[i][1]]
            nl = timesE[interval[i+1][0]]

            dist.append(nl-r)

        d_idx = {}
        for i in range(len(dist)):
            v = dist[i]
            d_idx[v] = i

        sorted_dist = sorted(dist)

        x1 = 0
        x2 = 0
        var_sum = 1e9

        for i in range(len(sorted_dist) - 1):
            # if i + 1 == 1:
            #     continue
            for j in np.arange(i+1,len(sorted_dist)):
                # if j - i == 1 or len(sorted_dist) - j - 1 == 1:
                #     continue

                seg1 = sorted_dist[0:i+1]
                seg2 = sorted_dist[i+1:j+1]
                seg3 = sorted_dist[j+1:]

                tvar1 = np.var(seg1)
                tvar2 = np.var(seg2)
                tvar3 = np.var(seg3)

                tvar_sum = tvar1 + tvar2 + tvar3
                if var_sum > tvar_sum:
                    var_sum = tvar_sum
                    x1 = i
                    x2 = j

        seg1 = sorted_dist[0:x1+1]
        seg2 = sorted_dist[x1+1:x2+1]
        seg3 = sorted_dist[x2+1:]

        mean1 = np.mean(seg1)
        mean2 = np.mean(seg2)
        mean3 = np.mean(seg3)
        tmean = mean1 + mean2

        missing_interval = []
        missing_maximum = []

        for seg_dist in seg3:
            id = d_idx[seg_dist]

            lidx = interval[id][1]
            ridx = interval[id+1][0]

            peaks_ = []
            for idx in peaks:
                if idx >= lidx and idx <= ridx:
                    peaks_.append(idx)


            sub_envelope = [envelope[x] for x in peaks_]
            # sub_var = np.var(sub_envelope)
            # left_bound = int(timesE[lidx] * framerate / downsampling_ratio)
            # right_bound = int(timesE[ridx] * framerate / downsampling_ratio)
            sub_wave_data = envelope[lidx:ridx+1]
            if len(sub_wave_data) == 0:
                continue

            # sub_minv = min(sub_wave_data)
            sub_minv = np.quantile(sub_wave_data, 0.25, interpolation='lower')  #
            sub_envelope_arr = np.array(sub_envelope)
            sub_envelope_arr = sub_envelope_arr - sub_minv

            if len(sub_envelope_arr) == 0:
                continue

            maxh = max(sub_envelope_arr)
            # medh = np.median(sub_envelope_arr)
            medh = np.quantile(sub_envelope_arr, 0.25, interpolation='lower')

            sub_var = np.var(sub_envelope_arr)
            # 判断是否有缺失
            if seg_dist > tmean * 0.8 and len(peaks_) > 1 and (sub_var > 0.01 and maxh > 2 * medh):

                sec = timesE[ridx] - timesE[lidx]

                t_interval,t_maximum,t_threshold1,t_threshold2 = search_time_door(timesE,envelope,peaks_,sec)

                missing_interval.extend(t_interval)
                missing_maximum.extend(t_maximum)

                interval, maximum = sort_interval(interval,missing_interval,maximum,missing_maximum)

        # drawTimeDoor(timesE, envelope, peaks, second, threshold1, threshold2, False, interval, maximum)


    interval, maximum = merge_interval(interval,maximum)
    # drawTimeDoor(timesE,envelope,peaks,second,threshold1,threshold2,False,interval,maximum)
    interval, maximum = filter_interval(timesE,interval,maximum,envelope,peaks)

    # drawTimeDoor(timesE,envelope,peaks,second,threshold1,threshold2,False,interval,maximum)


    return interval, maximum

    # --------------------------------------------------

def get_label(interval,timesE,wave_time,downsampling_ratio,framerate):
    dist = []
    for i in range(len(interval) - 1):
        r = timesE[interval[i][1]]
        nl = timesE[interval[i + 1][0]]

        dist.append(nl - r)

    maxd = -1e9
    maxid = 0
    for idx in range(len(dist)):
        curd = dist[idx]
        if curd > maxd:
            maxd = curd
            maxid = idx

    # 从最大间隔处开始打标签
    start_idx = maxid
    interval_label = {}
    front_num = 2 * (start_idx + 1)
    back_num = 2 * len(interval) - 2 * start_idx - 3

    # 计算前半段标签
    front_label = []
    v = 4
    for i in range(front_num):
        front_label.append(v)
        v -= 1
        if v == 0:
            v = 4

    if timesE[interval[0][0]] / (downsampling_ratio * 1.0 / framerate) > 0:
        front_label.append(v)

    front_label.reverse()

    # 计算后半段标签
    back_label = []
    v = 1
    for i in range(back_num):
        back_label.append(v)
        v += 1
        if v == 5:
            v = 1

    if timesE[interval[len(interval)-1][1]] / (downsampling_ratio * 1.0 / framerate) < len(wave_time) - 1:
        back_label.append(v)

    res_label = front_label
    res_label.extend(back_label)

    label_len = []
    last = -1
    for i in range(len(interval)):
        l = timesE[interval[i][0]] / (downsampling_ratio * 1.0 / framerate)
        r = timesE[interval[i][1]] / (downsampling_ratio * 1.0 / framerate)
        label_len.append(int(round(l-last-1)))
        label_len.append(int(round(r-l+1)))
        last = r

    end_idx = len(wave_time) - 1
    if last < end_idx:
        label_len.append(end_idx-last)

    labels = []

    for idx in range(len(res_label)):
        length = label_len[idx]
        label = res_label[idx]

        for j in range(int(length)):
            labels.append(label)

    return labels
